package org.budo.elasticsearch.jdbc.driver.statement;

import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.List;
import java.util.Set;

import org.budo.elasticsearch.jdbc.driver.BudoElasticsearchJdbcConnection;
import org.budo.elasticsearch.jdbc.driver.BudoElasticsearchJdbcPreparedStatement;
import org.budo.elasticsearch.jdbc.driver.resultset.SearchHitArrayResultSet;
import org.budo.jdbc.driver.BudoJdbcResultSetMetaData;
import org.budo.jdbc.driver.resultset.impl.ObjectArrayResultSet;
import org.elasticsearch.action.search.SearchRequestBuilder;
import org.elasticsearch.action.search.SearchType;
import org.elasticsearch.client.transport.TransportClient;
import org.elasticsearch.search.SearchHit;

/**
 * @author limingwei
 */
public class SelectStatement extends AbstractElasticsearchStatement {
    @Override
    public void execute(PreparedStatement preparedStatement) throws SQLException {
        if (this.isSelectCount()) { // SELECT COUNT(*) FROM table
            this.selectCount(preparedStatement);
            return;
        }

        this.selectList(preparedStatement);
    }

    private void selectCount(PreparedStatement preparedStatement) throws SQLException {
        BudoElasticsearchJdbcConnection connection = (BudoElasticsearchJdbcConnection) preparedStatement.getConnection();
        TransportClient transportClient = connection.getTransportClient();

        String elasticsearchQuery = this.elasticsearchQuery(preparedStatement, null); // 查询条件

        SearchRequestBuilder searchRequest = transportClient.prepareSearch() //
                .setIndices(new String[] { connection.getCatalog() }) // 库
                .setTypes(this.tableName()) // 表
                .setQuery(elasticsearchQuery) // 筛选条件
                .setSearchType(SearchType.QUERY_THEN_FETCH) // COUNT 类型
                .setSize(0);

        long count = searchRequest // 返回结果条数
                .get() //
                .getHits() //
                .totalHits();

        ObjectArrayResultSet resultSet = new ObjectArrayResultSet(new String[] { "COUNT(*)" }, new Object[] { count });

        ((BudoElasticsearchJdbcPreparedStatement) preparedStatement).setResultSet(resultSet);
    }

    private void selectList(PreparedStatement preparedStatement) throws SQLException {
        String[] selectColumnNames = this.getSelectItemColumnNames();
        SearchHit[] searchHits = this.executeSearchRequest(preparedStatement, selectColumnNames);

        // build resultSet
        SearchHitArrayResultSet resultSet = new SearchHitArrayResultSet();
        resultSet.setValueArray(searchHits);

        // 返回的列
        String[] resultSetColumnNames = this.resultSetColumnNames(selectColumnNames, searchHits);
        BudoJdbcResultSetMetaData metaData = new BudoJdbcResultSetMetaData(resultSetColumnNames);
        resultSet.setMetaData(metaData);

        ((BudoElasticsearchJdbcPreparedStatement) preparedStatement).setResultSet(resultSet);
    }

    private boolean isSelectCount() {
        return this.getSql() //
                .trim() //
                .toUpperCase() //
                .replace(" ", "") //
                .startsWith("SELECTCOUNT(*)FROM");
    }

    private String[] resultSetColumnNames(String[] selectColumnNames, SearchHit[] searchHits) {
        if (null == selectColumnNames || selectColumnNames.length < 1 || null == searchHits || searchHits.length < 1) {
            return selectColumnNames;
        }

        List<String> names = new ArrayList<String>();
        for (String selectColumnName : selectColumnNames) {
            if (selectColumnName.equals("*")) { // SELECT * 处理
                Set<String> keys = searchHits[0].getSource().keySet();
                names.addAll(keys);
                continue;
            }

            names.add(selectColumnName);
        }

        return names.toArray(new String[names.size()]);
    }
}